{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AvAmfzMk0YIu"
      },
      "source": [
        "# License\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\")\n",
        "```\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J-dZTKk9LM9K"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wipJHxX_LBtj"
      },
      "outputs": [],
      "source": [
        "# Uncomment to install the covid_vhh_design package\n",
        "\n",
        "# !pip install git+https://github.com/google-research/google-research.git#subdirectory=covid_vhh_design"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QYHih0_eoQ42"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "\n",
        "import collections\n",
        "import os\n",
        "\n",
        "import altair as alt\n",
        "from IPython import display\n",
        "import immutabledict\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import gridspec\n",
        "import seaborn as sns\n",
        "import statsmodels.stats.multitest\n",
        "import scipy\n",
        "import numpy as np\n",
        "import pandas as pd"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q9Nqo5nceiUq"
      },
      "outputs": [],
      "source": [
        "from covid_vhh_design import helper\n",
        "from covid_vhh_design import covid\n",
        "from covid_vhh_design import plotting\n",
        "from covid_vhh_design import utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zIKfbCo-S7_D"
      },
      "source": [
        "# Settings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o7LLtyX5TAaW"
      },
      "outputs": [],
      "source": [
        "plotting.update_rcparams()\n",
        "\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "pd.set_option('display.width', 200)\n",
        "pd.set_option('display.max_colwidth', None)\n",
        "pd.set_option('display.max_rows', 200)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QEBuidwkYDgT"
      },
      "source": [
        "# Constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CGqAFolTXP5A"
      },
      "outputs": [],
      "source": [
        "COV_WT = covid.COV_WT\n",
        "COV1_WT = covid.COV1_WT\n",
        "COV2_WT = covid.COV2_WT\n",
        "SHUFFLED = utils.SHUFFLED\n",
        "PARENT_NAME = covid.PARENT_NAME\n",
        "\n",
        "COLOR_PALETTE = plotting.PALETTE\n",
        "\n",
        "TARGET_NAME_MAPPING = covid.TARGET_SHORT_NAME_MAPPING\n",
        "TARGET_NAME_ORDER = tuple(list(covid.TARGET_SHORT_NAME_MAPPING.values()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ifz6aQYUTJRy"
      },
      "source": [
        "# Utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qOdRmdEBTKLc"
      },
      "outputs": [],
      "source": [
        "def concat_data_rounds(data_rounds):\n",
        "  return pd.concat(\n",
        "      [df.assign(round=i) for i, df in data_rounds.items()],\n",
        "      ignore_index=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rpp_Oof1sNsF"
      },
      "source": [
        "# Initialize global variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZXpxHofXsJnA"
      },
      "outputs": [],
      "source": [
        "Data = helper.Bunch()  # Data\n",
        "D = Data\n",
        "C = helper.Bunch()  # Computations\n",
        "T = helper.Bunch()  # Temporary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v6HgYkYDzgdZ"
      },
      "source": [
        "# Load and clean data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gaTLNs7YT-mF"
      },
      "source": [
        "## AlphaSeq data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "llVh6cDGovOL"
      },
      "outputs": [],
      "source": [
        "# @title Load data\n",
        "\n",
        "def load_alpha_seq_data():\n",
        "  \"\"\"Map 0-indexed round to corresponding library.\"\"\"\n",
        "  alpha_seq = {i: covid.load_alphaseq_data(i + 1) for i in range(3)}\n",
        "  return alpha_seq\n",
        "\n",
        "\n",
        "Data.raw_alpha_seq = load_alpha_seq_data()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y-yOpicC0hTs"
      },
      "outputs": [],
      "source": [
        "#@title Fix annotation mistakes\n",
        "def multiple_annotations_within_round(alphaseq_data):\n",
        "  df = pd.concat([df.assign(round=i) for i, df in alphaseq_data.items()],\n",
        "                 ignore_index=True)\n",
        "  return (df.groupby(['source_seq', 'round',\n",
        "                      'source_num_mutations'])['source_std_group'].agg([\n",
        "                          'nunique', 'unique'\n",
        "                      ]).query('nunique != 1')['unique'].reset_index())\n",
        "\n",
        "\n",
        "# This sequence is annotated both as `best_round1` and `baseline_r0`.\n",
        "_SPECIAL_SEQ = (\n",
        "    'QVQLQESGGGLVQAGGSLRLSCAASGFTFSEYAMGWFRQAPGKEREFVATISWSGRSTYYTD'\n",
        "    'SVKGRFTISRDNAKNTVYLQMNSLKPDDTAVYYCASAGLFTYVSEWDYDYDYWGQGTQVTVSS')\n",
        "\n",
        "mistaken_annotations = multiple_annotations_within_round(Data.raw_alpha_seq)\n",
        "mistaken_non_singles = mistaken_annotations[\n",
        "    mistaken_annotations['source_num_mutations'] != 1]\n",
        "assert len(mistaken_non_singles) == 1\n",
        "assert helper.get_unique_value(\n",
        "    mistaken_non_singles['source_seq']) == _SPECIAL_SEQ\n",
        "assert helper.get_unique_value(\n",
        "    mistaken_non_singles['round']) == 2\n",
        "\n",
        "# Correct annotation (singles are fixed in the `annotate` call).\n",
        "Data.raw_alpha_seq[2].loc[Data.raw_alpha_seq[2]['source_seq'] == _SPECIAL_SEQ,\n",
        "                          'source_group'] = 'baseline_r0'\n",
        "\n",
        "display.display(mistaken_annotations)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FAJ7JoR0U65D"
      },
      "outputs": [],
      "source": [
        "#@title Annotate \u0026 filter data\n",
        "%time Data.alpha_seq = covid.filter_alphaseq_data(covid.annotate_alphaseq_data(Data.raw_alpha_seq))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DK64ABMi2x72"
      },
      "outputs": [],
      "source": [
        "#@title Write updated annotations to disk\n",
        "def sequences_with_updated_annotations(alphaseq_data):\n",
        "  df = pd.concat([df.assign(round=i + 1) for i, df in alphaseq_data.items()],\n",
        "                 ignore_index=True)\n",
        "  updated_seqs = df[df['source_std_group'] !=\n",
        "                    df['new_source_std_group']]['source_seq'].unique()\n",
        "  df = df[df['source_seq'].isin(updated_seqs)]\n",
        "  df = df.groupby(['source_seq', 'round', 'new_source_std_group'\n",
        "                  ])['source_std_group'].first().reset_index()\n",
        "  return (df.pivot(\n",
        "      index=['source_seq', 'new_source_std_group'],\n",
        "      columns='round',\n",
        "      values='source_std_group'))\n",
        "\n",
        "\n",
        "updated_annos = sequences_with_updated_annotations(Data.alpha_seq).reset_index()\n",
        "display.display(updated_annos.head())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "teRR991LUuzN"
      },
      "source": [
        "## BLI data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xCkGAUtRUy20"
      },
      "outputs": [],
      "source": [
        "Data.bli_v1 = covid.load_df('bli_v1.csv')\n",
        "Data.bli_v2 = covid.load_df('bli_v2.csv')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4_J6GLhGzIOw"
      },
      "source": [
        "## Aggregate all sequence replicas"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S6FoU29Ktf7m"
      },
      "source": [
        "To compute the \"representative\" binding value of a sequence to a target, we\n",
        "1. Standardize by the experimental replicate (each experimental replicate has a mean of 0 and a standard deviation of 1 (`inf`s excluded).\n",
        "2. Standardize by the parent, so that the parent binding for each experimental replicate has a median of 0 and an IQR of 1.\n",
        "3. Median aggregate per sequence."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zgSXUFpq4dTW"
      },
      "outputs": [],
      "source": [
        "def get_metadata(alphaseq_data):\n",
        "  return {i: utils.get_metadata(alphaseq_data[i]) for i in alphaseq_data}\n",
        "\n",
        "\n",
        "def aggregate_affinities(alphaseq_data):\n",
        "  return {\n",
        "      i: utils.aggregate_affinities(alphaseq_data[i])\n",
        "      for i in alphaseq_data\n",
        "  }\n",
        "\n",
        "\n",
        "def compute_pvalues(alphaseq_data,\n",
        "                    alpha: float = 0.05,\n",
        "                    correction_method: str = 'bonferroni'):\n",
        "  return {\n",
        "      i: utils.compute_pvalues(\n",
        "          alphaseq_data[i],\n",
        "          alpha=alpha,\n",
        "          correction_method=correction_method,\n",
        "          min_counts=18) for i in alphaseq_data\n",
        "  }\n",
        "\n",
        "\n",
        "Data.metadata = get_metadata(Data.alpha_seq)\n",
        "Data.aggregated = aggregate_affinities(Data.alpha_seq)\n",
        "\n",
        "Data.r2_pvalues = utils.compute_pvalues(\n",
        "    Data.alpha_seq[2].query('target_name in @COV_WT'),\n",
        "    alpha=0.05,\n",
        "    correction_method='bonferroni',\n",
        "    min_replicas=18)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iqPUnMfy58WJ"
      },
      "outputs": [],
      "source": [
        "def join_aggregate_data_with_values(metadata, aggregated, pvalues):\n",
        "  return {\n",
        "      i: utils.join_aggregate_data_with_values(\n",
        "          metadata_df=metadata[i], agg_df=aggregated[i], pvalues_df=pvalues[i])\n",
        "      for i in range(3)\n",
        "  }\n",
        "\n",
        "\n",
        "Data.main = join_aggregate_data_with_values(\n",
        "    metadata=Data.metadata,\n",
        "    aggregated=Data.aggregated,\n",
        "    pvalues={\n",
        "        0: None,\n",
        "        1: None,\n",
        "        2: Data.r2_pvalues\n",
        "    })"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jAFjoTmNJ65g"
      },
      "outputs": [],
      "source": [
        "Data.unified = utils.aggregate_over_rounds(Data.alpha_seq).merge(\n",
        "    Data.r2_pvalues, how='left', on=['source_seq', 'target_name'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rd5bIQYA__we"
      },
      "source": [
        "# Standard deviation vs IQR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GUQjxof1kdXO"
      },
      "outputs": [],
      "source": [
        "def _compare_std_and_iqr(raw_df, target_name):\n",
        "  df = raw_df[raw_df['target_name'] == target_name].copy()\n",
        "  df = utils.standardize_experimental_replicates(df)\n",
        "  df = utils.extract_parent_df(df)\n",
        "\n",
        "  std_max_impute = np.std(helper.max_impute_inf(df['value']))\n",
        "  std_drop = np.std(helper.drop_inf_df(df, 'value')['value'])\n",
        "  return {\n",
        "      'std (max impute)': std_max_impute,\n",
        "      'std (drop inf)': std_drop,\n",
        "      'IQR': scipy.stats.iqr(df['value'])\n",
        "  }\n",
        "\n",
        "\n",
        "def compare_std_and_iqr(raw_data, target_name):\n",
        "\n",
        "  rows = []\n",
        "  for i in raw_data:\n",
        "    row = dict(round=i + 1)\n",
        "    row.update(_compare_std_and_iqr(raw_data[i], target_name))\n",
        "    rows.append(row)\n",
        "\n",
        "  return pd.DataFrame(rows).set_index('round')\n",
        "\n",
        "\n",
        "compare_std_and_iqr(Data.alpha_seq, COV2_WT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3AhDsJouGsCu"
      },
      "source": [
        "# AlphaSeq\n",
        "\n",
        "**Round 0**\n",
        "- 1 synthesis replica\n",
        "- 3 experimental replicas\n",
        "\n",
        "**Round 1**\n",
        "- 1 synthesis replica\n",
        "- 3 experimental replicas, but we drop the first one, leaving only `replica in [2, 3]`\n",
        "\n",
        "**Round 2**\n",
        "- 3 synthesis replicas\n",
        "- 6 experimental replicas\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "P4rkJWVcA76h"
      },
      "outputs": [],
      "source": [
        "#@title Experimental replicas {form-width: \"30%\"}\n",
        "\n",
        "source_replica = 1  #@param\n",
        "target_name = COV2_WT  #@param\n",
        "round = 0  #@param\n",
        "\n",
        "\n",
        "def _plot_experimental_replica_noise(ax, df, replica_x, replica_y):\n",
        "  df = df.pivot(\n",
        "      index=['source_seq', 'source_design'], values='value', columns='replica')\n",
        "  df.columns = [f'Exp. replica {c}' for c in df.columns]\n",
        "\n",
        "  plotting.plot_correlations(\n",
        "      ax,\n",
        "      df,\n",
        "      x_col=f'Exp. replica {replica_x}',\n",
        "      y_col=f'Exp. replica {replica_y}',\n",
        "      hue_col='source_design')\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(5, 5))\n",
        "_plot_experimental_replica_noise(\n",
        "    ax,\n",
        "    df= (\n",
        "        Data.alpha_seq[round]\n",
        "         .query('target_name == @target_name')\n",
        "         .query('source_replica == @source_replica')\n",
        "    ),\n",
        "    replica_x=2,\n",
        "    replica_y=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i5PFD4s3F0Bo"
      },
      "outputs": [],
      "source": [
        "#@title Synthesis replicas {form-width: \"30%\"}\n",
        "\n",
        "experimental_replica = 1  #@param\n",
        "target_name = COV2_WT  #@param\n",
        "round = 2  #@param\n",
        "\n",
        "\n",
        "def _plot_synthetic_replica_noise(ax, df, replica_x, replica_y):\n",
        "  df = df.pivot(\n",
        "      index=['source_seq', 'source_design'],\n",
        "      values='value',\n",
        "      columns='source_replica')\n",
        "  df.columns = [f'Syn. replica {int(c)}' for c in df.columns]\n",
        "\n",
        "  plotting.plot_correlations(\n",
        "      ax,\n",
        "      df,\n",
        "      x_col=f'Syn. replica {replica_x}',\n",
        "      y_col=f'Syn. replica {replica_y}',\n",
        "      hue_col='source_design')\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(5, 5))\n",
        "_plot_synthetic_replica_noise(\n",
        "    ax,\n",
        "    df=(\n",
        "        Data.alpha_seq[round]\n",
        "        .query('target_name == @target_name')\n",
        "        .query('replica == @experimental_replica')),\n",
        "    replica_x=1,\n",
        "    replica_y=2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PidCaLUhHUKf"
      },
      "outputs": [],
      "source": [
        "#@title By round {form-width: \"30%\"}\n",
        "round_x = 0  #@param\n",
        "round_y = 2  #@param\n",
        "target_name = COV2_WT  #@param\n",
        "\n",
        "\n",
        "def _plot_experimental_round_noise(ax, df, round_x, round_y):\n",
        "  df = df.pivot(\n",
        "      index=['source_seq', 'source_design'], values='value', columns='round')\n",
        "  df.columns = [f'Round {c+1}' for c in df.columns]\n",
        "  plotting.plot_correlations(\n",
        "      ax,\n",
        "      df,\n",
        "      x_col=f'Round {round_x + 1}',\n",
        "      y_col=f'Round {round_y + 1}',\n",
        "      hue_col='source_design')\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(5, 5))\n",
        "_plot_experimental_round_noise(\n",
        "    ax,\n",
        "    df=pd.concat([Data.main[x].assign(round=x) for x in range(3)],\n",
        "                 ignore_index=True).query('target_name==@target_name'),\n",
        "    round_x=round_x,\n",
        "    round_y=round_y)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jo-tQavYOJ8y"
      },
      "outputs": [],
      "source": [
        "#@title Source replicas {form-width: \"30%\"}\n",
        "\n",
        "target_name = COV2_WT  #@param\n",
        "\n",
        "\n",
        "def _plot_source_replica_noise(ax, df: pd.DataFrame, shrink=.8):\n",
        "  df = df[df['source_replica'].isin(\n",
        "      [1, 2, 3])].assign(source_replica=df['source_replica'].astype(int))\n",
        "  df = df.pivot_table(\n",
        "      index=['source_seq'],\n",
        "      values='value',\n",
        "      columns=['source_replica', 'replica'])\n",
        "\n",
        "  df = df.corr(method='spearman')\n",
        "  df.columns = df.columns.to_flat_index()\n",
        "  df.index = df.index.to_flat_index()\n",
        "\n",
        "  df.columns = df.columns.map(lambda x: f'{x[0]} - {x[1]}'\n",
        "                              if x[1] == 1 else x[1])\n",
        "  df.index = df.index.map(lambda x: f'{x[0]} - {x[1]}' if x[1] == 1 else x[1])\n",
        "\n",
        "  sns.heatmap(\n",
        "      df,\n",
        "      linewidth=.5,\n",
        "      vmin=0,\n",
        "      vmax=1,\n",
        "      cmap='rocket_r',\n",
        "      cbar_kws={'shrink': shrink},\n",
        "      ax=ax)\n",
        "  ax.set_aspect('equal', adjustable='box')\n",
        "  ax.set_xlabel('')\n",
        "  ax.set_ylabel('')\n",
        "  ax.xaxis.tick_top()  # x axis on top\n",
        "  ax.tick_params(axis='x', labelrotation=90, labelsize=14)\n",
        "  ax.tick_params(axis='y', labelrotation=0, labelsize=14)\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(6, 6))\n",
        "_plot_source_replica_noise(\n",
        "    ax,\n",
        "    df=Data.alpha_seq[2].query('target_name==@target_name'),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O0pA7Dzn5GDj"
      },
      "source": [
        "# $p$-value plots"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G0_FXHFy5Hnl"
      },
      "outputs": [],
      "source": [
        "def _plot_pvalue_vs_value(ax, df: pd.DataFrame, target_name=COV2_WT):\n",
        "  df = df[df['target_name'] == target_name]\n",
        "  sns.scatterplot(\n",
        "      data=df,\n",
        "      x='value',\n",
        "      y='pvalue',\n",
        "      ax=ax,\n",
        "      palette=plotting.PALETTE,\n",
        "      hue='source_design',\n",
        "  )\n",
        "  ax.set_xlabel('Normalized log KD')\n",
        "  ax.set_ylabel('Mann-Whitney U-test pvalue')\n",
        "  return ax\n",
        "\n",
        "\n",
        "def _plot_corrected_pvalue_vs_value(\n",
        "    ax,\n",
        "    df: pd.DataFrame,\n",
        "    correction_method: str,\n",
        "    target_name=COV2_WT,\n",
        "):\n",
        "  # We use pvalue == pvalue to drop NaN pvalues.\n",
        "  df = df.query('target_name == @target_name \u0026 pvalue == pvalue')\n",
        "  df = df.assign(\n",
        "      pvalue_corrected=statsmodels.stats.multitest.multipletests(\n",
        "          df['pvalue'], alpha=.05, method=correction_method)[1])\n",
        "\n",
        "  sns.scatterplot(\n",
        "      data=df,\n",
        "      x='value',\n",
        "      y='pvalue_corrected',\n",
        "      hue='source_design',\n",
        "      ax=ax,\n",
        "      palette=plotting.PALETTE)\n",
        "  ax.set_xlabel('Normalized log KD')\n",
        "  ax.set_ylabel(f'{correction_method.upper()}-corrected pvalue')\n",
        "  return ax\n",
        "\n",
        "\n",
        "f, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))\n",
        "_plot_pvalue_vs_value(ax1, Data.main[2])\n",
        "_plot_corrected_pvalue_vs_value(\n",
        "    ax2, Data.main[2], correction_method='bonferroni')\n",
        "\n",
        "for ax in (ax1, ax2):\n",
        "  ax.legend().set_visible(False)\n",
        "\n",
        "handles, labels = ax1.get_legend_handles_labels()\n",
        "f.legend(handles, labels, ncol=4, bbox_to_anchor=(.7, 1.0), frameon=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gv0jQE2vSOdY"
      },
      "source": [
        "# Percentage of Infs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sz5H_bdOScoL"
      },
      "outputs": [],
      "source": [
        "def _compute_percentage_infs(values):\n",
        "  is_infs = np.isinf(values)\n",
        "  return (\n",
        "      pd.DataFrame(dict(\n",
        "        num_infs=is_infs.sum(),\n",
        "        frac_infs=is_infs.mean(),\n",
        "        num_samples=len(values)),\n",
        "        index=[0])\n",
        "      .assign(\n",
        "          percentage_infs=lambda df: df['frac_infs'] * 100))\n",
        "\n",
        "\n",
        "def compute_percentage_infs(df, group_vars):\n",
        "  group_vars = helper.to_list(group_vars)\n",
        "  return (\n",
        "      df.groupby(group_vars)['value']\n",
        "      .apply(_compute_percentage_infs)\n",
        "      .reset_index(group_vars))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WoWzhtrySeid"
      },
      "source": [
        "## All rounds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SZLA8UOFWcD3"
      },
      "outputs": [],
      "source": [
        "def _plot_percentage_infs_rounds(ax, data):\n",
        "  df = (\n",
        "      helper.map_columns(\n",
        "          data,\n",
        "          target_name=covid.TARGET_SHORT_NAME_MAPPING,\n",
        "          round=covid.ROUND_MAPPING)\n",
        "  )\n",
        "\n",
        "  sns.barplot(\n",
        "      ax=ax,\n",
        "      data=df,\n",
        "      x='round',\n",
        "      order=covid.ROUND_MAPPING.values(),\n",
        "      y='percentage_infs',\n",
        "      hue='target_name',\n",
        "      palette=plotting.PALETTE,\n",
        "  )\n",
        "  plotting.rotate_xlabels(ax, 30)\n",
        "  ax.set_ylabel('% of non-binders')\n",
        "  ax.legend(title='Target')\n",
        "  return ax\n",
        "\n",
        "\n",
        "C.per_infs_rounds = compute_percentage_infs(\n",
        "    concat_data_rounds(Data.alpha_seq)\n",
        "    .query('target_name in @COV_WT'),\n",
        "    group_vars=['target_name', 'round'])\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(6, 4))\n",
        "_plot_percentage_infs_rounds(ax, C.per_infs_rounds)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P7_QRcyHXnCf"
      },
      "outputs": [],
      "source": [
        "def _plot_percentage_infs_by_model(ax, data, model_mapping):\n",
        "  df = (\n",
        "      data.query('source_group in @model_mapping').pipe(\n",
        "          helper.map_columns,\n",
        "          target_name=covid.TARGET_SHORT_NAME_MAPPING,\n",
        "          source_group=model_mapping))\n",
        "  sns.barplot(\n",
        "      ax=ax,\n",
        "      data=df,\n",
        "      x='source_group',\n",
        "      order=model_mapping.values(),\n",
        "      y='percentage_infs',\n",
        "      hue='target_name',\n",
        "      palette=plotting.PALETTE)\n",
        "  plotting.rotate_xlabels(ax, 30)\n",
        "  ax.set_ylabel('% of non-binders')\n",
        "  ax.legend(\n",
        "      bbox_to_anchor=(0.5, 1.1),\n",
        "      ncol=2,\n",
        "      loc='center',\n",
        "      frameon=False,\n",
        "  )\n",
        "\n",
        "  return ax"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CC0iAtmYlRw0"
      },
      "outputs": [],
      "source": [
        "#@title Round 1\n",
        "\n",
        "C.per_infs_r0 = compute_percentage_infs(\n",
        "    Data.alpha_seq[0].query('target_name in @COV_WT'),\n",
        "    group_vars=['target_name', 'source_group'])\n",
        "\n",
        "_, ax = plt.subplots(figsize=(10, 4))\n",
        "_plot_percentage_infs_by_model(\n",
        "    ax,\n",
        "    C.per_infs_r0,\n",
        "    model_mapping=covid.ROUND0_SOURCE_GROUP_MAPPING)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IuZe6Xv03gp2"
      },
      "outputs": [],
      "source": [
        "x = 1.058224\n",
        "y = 2.128489\n",
        "\n",
        "print(f'Pearson score={x:.2f}\\np-value={y:.2e}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g_fVxr-GlT18"
      },
      "outputs": [],
      "source": [
        "#@title Round 2\n",
        "\n",
        "C.per_infs_r1 = compute_percentage_infs(\n",
        "    Data.alpha_seq[1].query('target_name in @COV_WT'),\n",
        "    group_vars=['target_name', 'source_group'])\n",
        "\n",
        "_, ax = plt.subplots(figsize=(10, 4))\n",
        "_plot_percentage_infs_by_model(\n",
        "    ax, C.per_infs_r1, model_mapping=covid.ROUND1_SOURCE_GROUP_MAPPING)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xt8JVgGkZ5uN"
      },
      "outputs": [],
      "source": [
        "#@title Round 3\n",
        "\n",
        "C.per_infs_r2 = compute_percentage_infs(\n",
        "    Data.alpha_seq[2].query('target_name in @COV_WT'),\n",
        "    group_vars=['target_name', 'source_group'])\n",
        "\n",
        "_, ax = plt.subplots(figsize=(10, 4))\n",
        "_plot_percentage_infs_by_model(\n",
        "    ax,\n",
        "    C.per_infs_r2,\n",
        "    model_mapping=covid.ROUND2_SOURCE_GROUP_MAPPING)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r9U1IWlUgqSl"
      },
      "source": [
        "## By distance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7L-4pkMnheJz"
      },
      "outputs": [],
      "source": [
        "def _plot_percentage_infs_by_distance(ax, data):\n",
        "  df = (\n",
        "      data.query('source_num_mutations \u003e 0')\n",
        "      .pipe(\n",
        "          helper.map_columns,\n",
        "          target_name=covid.TARGET_SHORT_NAME_MAPPING)\n",
        "      .assign(\n",
        "          source_num_mutations=lambda df: df['source_num_mutations'].astype('int'))\n",
        "  )\n",
        "  ax = sns.barplot(\n",
        "      ax=ax,\n",
        "      data=df,\n",
        "      x='source_num_mutations',\n",
        "      y='percentage_infs',\n",
        "      hue='target_name',\n",
        "      palette=plotting.PALETTE)\n",
        "  ax.set_xlabel('Number of mutations from VHH-72')\n",
        "  ax.set_ylabel('% of non-binders')\n",
        "  ax.legend(\n",
        "      bbox_to_anchor=(0.5, 1.1),\n",
        "      ncol=2,\n",
        "      loc='center',\n",
        "      frameon=False,\n",
        "  )\n",
        "  ax.set_xticklabels(ax.get_xticklabels(), ha='center')\n",
        "\n",
        "  return ax\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jp_YQ5HjgtLV"
      },
      "outputs": [],
      "source": [
        "#@title Round 1\n",
        "C.per_infs_r0_dist = compute_percentage_infs(\n",
        "    Data.alpha_seq[0].query('target_name in @COV_WT'),\n",
        "    group_vars=['target_name', 'source_num_mutations'])\n",
        "\n",
        "_, ax = plt.subplots(figsize=(5.5, 4))\n",
        "_plot_percentage_infs_by_distance(ax, C.per_infs_r0_dist)\n",
        "ax.legend().set_visible(False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BgwhzvmGhRdk"
      },
      "outputs": [],
      "source": [
        "#@title Round 2\n",
        "C.per_infs_r1_dist = compute_percentage_infs(\n",
        "    Data.alpha_seq[1].query(\n",
        "        'target_name in @COV_WT and source_design != @SHUFFLED'),\n",
        "    group_vars=['target_name', 'source_num_mutations'])\n",
        "\n",
        "_, ax = plt.subplots(figsize=(5.5, 4))\n",
        "_plot_percentage_infs_by_distance(ax, C.per_infs_r1_dist)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N6Khjlt5hlDb"
      },
      "outputs": [],
      "source": [
        "#@title Round 3\n",
        "C.per_infs_r2_dist = compute_percentage_infs(\n",
        "    Data.alpha_seq[2].query(\n",
        "        'target_name in @COV_WT and source_design != @SHUFFLED'),\n",
        "    group_vars=['target_name', 'source_num_mutations'])\n",
        "\n",
        "_, ax = plt.subplots(figsize=(5.5, 4))\n",
        "_plot_percentage_infs_by_distance(ax, C.per_infs_r2_dist)\n",
        "ax.legend().set_visible(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oWSRjDdEMTUG"
      },
      "source": [
        "# log KD plots"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0zX2RnFCaHDf"
      },
      "source": [
        "## All rounds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rgrw27mHlan_"
      },
      "outputs": [],
      "source": [
        "def _plot_log_kd_by_round(ax, df, n=100, model_as_design: bool = False):\n",
        "  df = df.query('target_name == @COV2_WT and source_design != @SHUFFLED')\n",
        "  round2_df = df.query('round == 2 | source_num_mutations == 0')\n",
        "  if model_as_design:\n",
        "    round2_df = plotting.replace_ml_design_by_model(round2_df)\n",
        "\n",
        "  # We only want new hues for the last round (assuming `model_as_design=True`),\n",
        "  # so we're going to have to be sneaky.\n",
        "  # We are going to plot to boxplots. The first one plots the first two rounds,\n",
        "  # with hues 'Baseline' and 'ML', in the simple call below.\n",
        "  ax = plotting.plot_log_kd(\n",
        "      ax,\n",
        "      agg_df=df.query('round \u003c 2'),\n",
        "      max_impute_inf=True,\n",
        "      num_top_seqs=n,\n",
        "      x_col='round',\n",
        "      hue_col='source_design',\n",
        "      show_iqr=False,\n",
        "  )\n",
        "\n",
        "  # The second plotting call does some manipulations to get the alignment to\n",
        "  # work out: we only plot round 3, with its own hues ('Baseline', 'LGB', 'CNN')\n",
        "  # but to get theses boxes correctly aligned on the x-axis, we \"mock\" Baseline\n",
        "  # values (with some absurdly large numbers that won't show up after yaxis\n",
        "  # rescaling) for rounds 1 and 2.\n",
        "  mocked_prev_round_values = pd.DataFrame(\n",
        "      dict(\n",
        "          round=[0, 1] * n,\n",
        "          source_design='Baseline',\n",
        "          value=100,\n",
        "          source_num_mutations=1,\n",
        "      )\n",
        "  )\n",
        "  plotting.plot_log_kd(\n",
        "      ax,\n",
        "      agg_df=pd.concat([round2_df, mocked_prev_round_values]),\n",
        "      max_impute_inf=True,\n",
        "      num_top_seqs=n,\n",
        "      x_col='round',\n",
        "      hue_col='source_design',\n",
        "      show_iqr=True,\n",
        "  )\n",
        "  ax.set_ylim(-4.3, 1.3)\n",
        "\n",
        "  ax.set_title(\n",
        "      f'Log KDs of each round\\'s top {n} sequences against CoV-2', y=-0.2\n",
        "  )\n",
        "  ax.set_ylabel('Normalized log KD against CoV-2')\n",
        "\n",
        "  # Finally, some of the labels are duplicated in the legend (because of two\n",
        "  # plotting calls), so we need to deduplicate.\n",
        "  handles, labels = ax.get_legend_handles_labels()\n",
        "  labels_to_handles = dict(zip(labels, handles))\n",
        "  handles = list(labels_to_handles.values())\n",
        "  labels = list(labels_to_handles.keys())\n",
        "  ax.legend(\n",
        "      handles=handles,\n",
        "      labels=labels,\n",
        "      bbox_to_anchor=(0.5, 1.1),\n",
        "      ncol=5,\n",
        "      loc='center',\n",
        "      frameon=False,\n",
        "  )\n",
        "  return ax\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(10, 5))\n",
        "ax = _plot_log_kd_by_round(ax, Data.unified, n=100, model_as_design=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TO8y9wXhbKOM"
      },
      "source": [
        "## All rounds by distance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vh5P8OPGlemh"
      },
      "outputs": [],
      "source": [
        "def _plot_log_kd_by_mutations(ax, df, n=100):\n",
        "  df = df.query(\n",
        "      'target_name == @COV2_WT and source_design != @SHUFFLED and (round == 2 or source_num_mutations == 0)'\n",
        "  )\n",
        "\n",
        "  plotting.plot_log_kd(\n",
        "      ax,\n",
        "      agg_df=df,\n",
        "      max_impute_inf=True,\n",
        "      num_top_seqs=n,\n",
        "      x_col='source_num_mutations',\n",
        "      hue_col='source_design',\n",
        "      show_iqr=True,\n",
        "  )\n",
        "\n",
        "  ax.set_title(f'Log KDs of top {n} round 3 sequences against CoV-2', y=-.3)\n",
        "  ax.set_ylabel('Normalized log KD against CoV-2')\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(12, 5))\n",
        "_plot_log_kd_by_mutations(ax, Data.unified)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dCWdtJm_aZvD"
      },
      "source": [
        "## All rounds separately by distance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7bSQVwdKlgis"
      },
      "outputs": [],
      "source": [
        "def _plot_all_rounds_by_mutations(ax, data, *, include_shuffled, ml_only):\n",
        "  df = data.query('target_name == @COV2_WT').assign(\n",
        "      round=data['round'].apply(lambda x: f'Round {x+1}')\n",
        "  )\n",
        "  if ml_only:\n",
        "    df = df.query(\n",
        "        'source_design == @utils.ML or source_num_mutations == 0'\n",
        "    )\n",
        "\n",
        "  elif not include_shuffled:\n",
        "    df = df.query('source_design != @SHUFFLED')\n",
        "\n",
        "  ax = plotting.plot_log_kd(\n",
        "      ax=ax,\n",
        "      agg_df=df,\n",
        "      x_col='source_num_mutations',\n",
        "      hue_col='round',\n",
        "      palette='Dark2',\n",
        "      max_impute_inf=True,\n",
        "      show_iqr=True,\n",
        "  )\n",
        "  ax.legend(loc='center', bbox_to_anchor=(0.5, 1.1), ncol=4, frameon=False)\n",
        "  ax.set_ylabel('Normalized log KD against CoV-2')\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(12, 5))\n",
        "_plot_all_rounds_by_mutations(\n",
        "    ax, Data.unified, include_shuffled=False, ml_only=True\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GfTPtfyykrPy"
      },
      "source": [
        "## By models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZCQJ--twljVY"
      },
      "outputs": [],
      "source": [
        "def _plot_log_kd_by_model(ax, df: pd.DataFrame, model_mapping: dict[str, str]):\n",
        "  df = df[df['target_short_name'].isin(['SARS-CoV-1', 'SARS-CoV-2'])].copy()\n",
        "\n",
        "  # TODO: The text should note what mutants are and that we drop them here.\n",
        "  df = df[df['source_group'] != 'mutant']\n",
        "  df = df.assign(source_model=df['source_group'].map(model_mapping))\n",
        "\n",
        "  # We take the median because there are slightly different values for the\n",
        "  # parent (all close to zero) for both targets.\n",
        "  assert utils.extract_parent_df(df)['value'].values.max() \u003c 1e-10\n",
        "  df.loc[df['source_num_mutations'] == 0, 'value'] = 0.\n",
        "\n",
        "  plotting.plot_log_kd(\n",
        "      ax,\n",
        "      agg_df=df,\n",
        "      max_impute_inf=True,\n",
        "      x_col='source_model',\n",
        "      order=model_mapping.values(),\n",
        "      hue_col='target_short_name',\n",
        "      show_iqr=True,\n",
        "  )\n",
        "  ax.set_xlabel('')\n",
        "  ax.set_xticklabels(labels=ax.get_xticklabels(), rotation=30, ha='right')\n",
        "  ax.legend(bbox_to_anchor=(.5, 1.1), loc='center', ncol=3, frameon=False)\n",
        "  return ax\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9lCi2lalllIK"
      },
      "outputs": [],
      "source": [
        "#@title Round 1\n",
        "_, ax = plt.subplots(figsize=(10, 4))\n",
        "_plot_log_kd_by_model(ax, Data.main[0], model_mapping=covid.ROUND0_SOURCE_GROUP_MAPPING)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7cURNjMZlROp"
      },
      "outputs": [],
      "source": [
        "#@title Round 2\n",
        "_, ax = plt.subplots(figsize=(10, 4))\n",
        "_plot_log_kd_by_model(ax, Data.main[1], model_mapping=covid.ROUND1_SOURCE_GROUP_MAPPING)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6xvmHEZ2lZRW"
      },
      "outputs": [],
      "source": [
        "#@title Round 3\n",
        "_, ax = plt.subplots(figsize=(6, 4))\n",
        "_plot_log_kd_by_model(\n",
        "    ax, Data.main[2], model_mapping=covid.ROUND2_SOURCE_GROUP_MAPPING)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4VDt-UEos5TM"
      },
      "source": [
        "## By distance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2kRzGlWsuRWU"
      },
      "outputs": [],
      "source": [
        "def _plot_round_by_mutations(ax, df: pd.DataFrame):\n",
        "  df = df[df['target_short_name'].isin(['SARS-CoV-1', 'SARS-CoV-2'])].copy()\n",
        "\n",
        "  # Since we plot multiple targets, there may be slightly different values of\n",
        "  # the parent binding.\n",
        "  assert utils.extract_parent_df(df)['value'].values.max() \u003c 1e-10\n",
        "  df.loc[df['source_num_mutations'] == 0, 'value'] = 0.\n",
        "\n",
        "  plotting.plot_log_kd(\n",
        "      ax=ax,\n",
        "      agg_df=df,\n",
        "      x_col='source_num_mutations',\n",
        "      hue_col='target_short_name',\n",
        "      max_impute_inf=True,\n",
        "      show_iqr=True,\n",
        "  )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GOauYCfyuqOL"
      },
      "outputs": [],
      "source": [
        "#@title Round 1\n",
        "_, ax = plt.subplots(figsize=(7, 5))\n",
        "_plot_round_by_mutations(ax, Data.main[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-9O-tbb2tH1K"
      },
      "outputs": [],
      "source": [
        "#@title Round 2\n",
        "_, ax = plt.subplots(figsize=(6, 4))\n",
        "_plot_round_by_mutations(ax, Data.main[1].query('source_design != @SHUFFLED'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7LlJ-2O5thwP"
      },
      "outputs": [],
      "source": [
        "#@title Round 3\n",
        "_, ax = plt.subplots(figsize=(10, 5))\n",
        "_plot_round_by_mutations(ax, Data.main[2].query('source_design != @SHUFFLED'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DHYq2g4Hv1fQ"
      },
      "source": [
        "# Supplemental figures for each round\n",
        "\n",
        "Using gridspec to avoid manually dealing with alignment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "azupFQsv3iYX"
      },
      "outputs": [],
      "source": [
        "def supplemental_round_analysis_figure(alphaseq_data, agg_data, round_idx: int):\n",
        "  width_ratios = [3, 2] if round_idx in [0, 1] else [2, 3]\n",
        "\n",
        "  model_mappings = {\n",
        "      0:\n",
        "          covid.ROUND0_SOURCE_GROUP_MAPPING,\n",
        "      1:\n",
        "          covid.ROUND1_SOURCE_GROUP_MAPPING,\n",
        "      2:\n",
        "          covid.ROUND2_SOURCE_GROUP_MAPPING | {\n",
        "              'baseline_r0': 'Recomb. \u0026 \\nsingles (R1)',\n",
        "              'baseline_r1': 'Recomb. \u0026 \\nsingles (R2)',\n",
        "          },\n",
        "  }\n",
        "\n",
        "  fig = plt.figure(figsize=(20, 16))\n",
        "  legend_spec = gridspec.GridSpec(\n",
        "      3,\n",
        "      2,\n",
        "      top=.8,\n",
        "      height_ratios=[.15, 1, 1],\n",
        "  )\n",
        "  figure_spec = gridspec.GridSpec(\n",
        "      3,\n",
        "      2,\n",
        "      width_ratios=width_ratios,\n",
        "      hspace=0.5,\n",
        "      height_ratios=[.15, 1, 1],\n",
        "  )\n",
        "\n",
        "  legend_ax = fig.add_subplot(legend_spec[0, :])\n",
        "  plotting.make_target_legend(ax=legend_ax, with_parent=True)\n",
        "\n",
        "  agg_df = agg_data[round_idx]\n",
        "  alphaseq_df = alphaseq_data[round_idx]\n",
        "\n",
        "  ax1 = fig.add_subplot(figure_spec[1, 0])\n",
        "  _plot_log_kd_by_model(ax1, agg_df, model_mapping=model_mappings[round_idx])\n",
        "  ax1.legend().set_visible(False)\n",
        "\n",
        "  ax2 = fig.add_subplot(figure_spec[1, 1])\n",
        "  _plot_round_by_mutations(\n",
        "      ax2, agg_df.query('source_design != @SHUFFLED'))\n",
        "  ax2.legend().set_visible(False)\n",
        "\n",
        "  ax3 = fig.add_subplot(figure_spec[2, 0])\n",
        "  per_infs = compute_percentage_infs(\n",
        "      alphaseq_df.query('target_name in @COV_WT'),\n",
        "      group_vars=['target_name', 'source_group'])\n",
        "  _plot_percentage_infs_by_model(\n",
        "      ax3, per_infs, model_mapping=model_mappings[round_idx])\n",
        "  ax3.legend().set_visible(False)\n",
        "\n",
        "  ax4 = fig.add_subplot(figure_spec[2, 1])\n",
        "  per_infs_dist = compute_percentage_infs(\n",
        "      alphaseq_df.query(\n",
        "          'target_name in @COV_WT and source_design != @SHUFFLED'),\n",
        "      group_vars=['target_name', 'source_num_mutations'])\n",
        "  _plot_percentage_infs_by_distance(ax4, per_infs_dist)\n",
        "  ax4.legend().set_visible(False)\n",
        "\n",
        "\n",
        "supplemental_round_analysis_figure(Data.alpha_seq, Data.main, round_idx=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Te24VHZYDZR1"
      },
      "outputs": [],
      "source": [
        "supplemental_round_analysis_figure(Data.alpha_seq, Data.main, round_idx=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tviszye7Dbhu"
      },
      "outputs": [],
      "source": [
        "supplemental_round_analysis_figure(Data.alpha_seq, Data.main, round_idx=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0YgxyoF1S57n"
      },
      "source": [
        "# Hit rate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3pN7kgXldSmo"
      },
      "outputs": [],
      "source": [
        "#@title By distance\n",
        "\n",
        "def _plot_hits_by_mutations(ax, df, model_as_hue: False):\n",
        "  df = df.query(\n",
        "      'target_name == @COV2_WT and pvalue_corrected == pvalue_corrected '\n",
        "      'and (round == 2 or source_num_mutations == 0)'\n",
        "      )\n",
        "\n",
        "  if model_as_hue:\n",
        "    df = plotting.replace_ml_design_by_model(df)\n",
        "\n",
        "  plotting.plot_hit_rate(\n",
        "      ax=ax,\n",
        "      agg_df=df,\n",
        "      how='pvalue',\n",
        "      x_col='source_num_mutations',\n",
        "      hue_col='source_design',\n",
        "      show_counts=True,\n",
        "  )\n",
        "  ax.set_title(\n",
        "      r'Hits against SARS-CoV-2 (p $\\leq$ 0.05) in round 3', y=-.3)\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(12, 5))\n",
        "_plot_hits_by_mutations(ax, Data.unified, model_as_hue=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5m8JL320A8LH"
      },
      "outputs": [],
      "source": [
        "#@title Comparison by pvalue correction\n",
        "f, axes = plt.subplots(1, 3, figsize=(30, 5))\n",
        "\n",
        "df = Data.unified.query(\n",
        "    'target_name == @COV2_WT and pvalue_corrected == pvalue_corrected '\n",
        "    'and (round == 2 or source_num_mutations == 0)')\n",
        "df = df.assign(\n",
        "    fdr_by=statsmodels.stats.multitest.multipletests(\n",
        "        df['pvalue'], alpha=.05, method='fdr_by')[1],\n",
        "    fdr_bh=statsmodels.stats.multitest.multipletests(\n",
        "        df['pvalue'], alpha=.05, method='fdr_bh')[1],\n",
        "    bonferroni=statsmodels.stats.multitest.multipletests(\n",
        "        df['pvalue'], alpha=.05, method='bonferroni')[1])\n",
        "\n",
        "title = 'Round 3'\n",
        "\n",
        "for ax, method in zip(axes, ['fdr_by', 'fdr_bh', 'bonferroni']):\n",
        "  plotting.plot_hit_rate(\n",
        "      ax=ax,\n",
        "      agg_df=df.assign(pvalue_corrected=df[method]),\n",
        "      how='pvalue',\n",
        "      x_col='source_num_mutations',\n",
        "      hue_col='source_design',\n",
        "      show_counts=True)\n",
        "  ax.set_ylabel(method.upper())\n",
        "\n",
        "f.suptitle(title, y=-.03)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u5fafGfzxo6o"
      },
      "source": [
        "## Hits in last round, aggregated"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W8qDcXYdxjXw"
      },
      "outputs": [],
      "source": [
        "def analyze_hits(df):\n",
        "  df = df.assign(significant_hit=Data.unified['pvalue_corrected'] \u003c= .05)\n",
        "  df = df[(df['round'] == 2) | (df['source_num_mutations'] == 0)]\n",
        "  df = df[df.target_name == COV2_WT]\n",
        "\n",
        "  display.display(\n",
        "      df.groupby('source_design')['significant_hit'].agg(\n",
        "          ['sum', 'count', 'mean']))\n",
        "\n",
        "\n",
        "analyze_hits(Data.unified)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ifOPikqFmA2E"
      },
      "outputs": [],
      "source": [
        "#@title By targets\n",
        "\n",
        "def _cov1_hits_by_iqr(row):\n",
        "  # Mock a \"significant hit\" by setting a pvalue \u003c .05.\n",
        "  if row['target_name'] == COV1_WT:\n",
        "    return 0 if row['value'] \u003c= 0 else 1\n",
        "\n",
        "  return row['pvalue_corrected']\n",
        "\n",
        "\n",
        "def add_joint_targets(df, how: str):\n",
        "  \"\"\"Adds rows corresponding to hits to both CoV1 and CoV2.\"\"\"\n",
        "  df = df[df['target_name'].isin([COV1_WT, COV2_WT])].copy()\n",
        "  df['target_name'] = df['target_name'].map(\n",
        "      covid.TARGET_SHORT_NAME_MAPPING)\n",
        "\n",
        "\n",
        "  if how == 'pvalue':\n",
        "    measurement_col = 'pvalue_corrected'\n",
        "  elif how == 'iqr':\n",
        "    measurement_col='value'\n",
        "  else:\n",
        "    raise ValueError('`how` argument must be either \"pvalue\" or \"iqr\".')\n",
        "\n",
        "  worst_binding = (\n",
        "      df.groupby([\n",
        "          'source_seq', 'source_design', 'source_num_mutations'\n",
        "      ])[measurement_col].max().reset_index().assign(target_name='CoV1 \u0026 CoV2'))\n",
        "  return pd.concat([df, worst_binding], ignore_index=True)\n",
        "\n",
        "\n",
        "def _plot_hits_by_targets(ax, df, how, cov1_as_iqr: bool):\n",
        "  df = df[(df['round'] == 2) | (df['source_num_mutations'] == 0)]\n",
        "  if cov1_as_iqr:\n",
        "    df = df.assign(pvalue_corrected=df.apply(_cov1_hits_by_iqr, axis=1))\n",
        "\n",
        "  df = add_joint_targets(plotting.replace_ml_design_by_model(df), how=how)\n",
        "\n",
        "  plotting.plot_hit_rate(\n",
        "      ax,\n",
        "      agg_df=df,\n",
        "      how=how,\n",
        "      x_col='target_name',\n",
        "      hue_col='source_design',\n",
        "      order=('SARS-CoV-1', 'SARS-CoV-2', 'CoV1 \u0026 CoV2'),\n",
        "      show_counts=True)\n",
        "  ax.set_title('Hits by target in round 3.', y=-.2)\n",
        "  return df\n",
        "\n",
        "\n",
        "_, ax = plt.subplots(figsize=(7, 5.5))\n",
        "df = _plot_hits_by_targets(ax, Data.unified, how='pvalue', cov1_as_iqr=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7NjXvnE4mnOl"
      },
      "outputs": [],
      "source": [
        "# @title Hits analysis\n",
        "\n",
        "\n",
        "def describe_mutations(seq, ipos_to_pos, ipos_to_region):\n",
        "  mutated_pos = []\n",
        "  mutated_regions = set()\n",
        "  for i, (mut, wt) in enumerate(zip(seq, covid.PARENT_SEQ)):\n",
        "    if mut != wt:\n",
        "      mutated_pos.append(ipos_to_pos[i])\n",
        "      mutated_regions.add(ipos_to_region[i])\n",
        "  return mutated_pos, sorted(mutated_regions)\n",
        "\n",
        "\n",
        "def describe_hits(df, target_name):\n",
        "  df = df[df['target_name'] == target_name].copy()\n",
        "  parent_seq = covid.load_aligned_parent_seq(offset_ipos=0)\n",
        "  ipos_to_pos = dict(zip(parent_seq['ipos'], parent_seq['pos']))\n",
        "  ipos_to_region = dict(zip(parent_seq['ipos'], parent_seq['region']))\n",
        "\n",
        "  parent_binding = helper.get_unique_value(\n",
        "      utils.extract_parent_df(df)['value']\n",
        "  )\n",
        "  hits = df[df['value'] \u003c= parent_binding - 1].copy()\n",
        "  hits['pos'], hits['regions'] = zip(\n",
        "      *hits['source_seq'].apply(\n",
        "          describe_mutations,\n",
        "          ipos_to_pos=ipos_to_pos,\n",
        "          ipos_to_region=ipos_to_region,\n",
        "      )\n",
        "  )\n",
        "\n",
        "  return hits[[\n",
        "      'source_seq',\n",
        "      'source_model',\n",
        "      'source_group',\n",
        "      'pos',\n",
        "      'regions',\n",
        "      'value',\n",
        "  ]].sort_values(by='value')\n",
        "\n",
        "\n",
        "design_round = 0\n",
        "print(\n",
        "    f'Sequences with \u003e= 1 IQR improvement over parent in round {design_round}'\n",
        ")\n",
        "describe_hits(Data.main[design_round], target_name=COV2_WT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Chgfa5EFeWb"
      },
      "source": [
        "# Diversity"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "swTuF4yMF2NK"
      },
      "source": [
        "### t-SNE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "adcrD64MJvc9"
      },
      "outputs": [],
      "source": [
        "Data.tsne = utils.compute_tsne_embedding(\n",
        "    Data.unified[Data.unified['target_name'] == COV2_WT])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XeKpDsEqlDrY"
      },
      "outputs": [],
      "source": [
        "Data.tsne['bli_v1'] = Data.tsne['source_seq'].apply(\n",
        "    lambda x: x in Data.bli_v1.source_seq.unique())\n",
        "Data.tsne['bli_v2'] = Data.tsne['source_seq'].apply(\n",
        "    lambda x: x in Data.bli_v2.source_seq.unique())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ERbXnYnwm2EL"
      },
      "outputs": [],
      "source": [
        "_, ax = plt.subplots(figsize=(7, 5))\n",
        "plotting.plot_tsne_by_round(ax=ax, agg_df=Data.tsne, round_idx=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "npsyyogAGCPh"
      },
      "outputs": [],
      "source": [
        "def make_tsne_figure():\n",
        "\n",
        "  spec = gridspec.GridSpec(\n",
        "      ncols=3,\n",
        "      nrows=2,\n",
        "      width_ratios=[1, 1, 1],\n",
        "      wspace=0.25,\n",
        "      hspace=0.2,\n",
        "      height_ratios=[.05, 1])\n",
        "\n",
        "  fig = plt.figure(figsize=(20, 5.5))\n",
        "  bli_legend_ax = fig.add_subplot(spec[0, :])\n",
        "  plotting.make_tsne_legend(ax=bli_legend_ax, with_initial_bli=True)\n",
        "\n",
        "  # Fig A\n",
        "  ax1 = fig.add_subplot(spec[1, 0])\n",
        "  plotting.plot_tsne_by_round(ax1, Data.tsne, round_idx=0)\n",
        "  ax1.set_title('Round 1', y=-.3)\n",
        "\n",
        "  ax2 = fig.add_subplot(spec[1, 1])\n",
        "  plotting.plot_tsne_by_round(ax2, Data.tsne, round_idx=1)\n",
        "  ax2.set_title('Round 2', y=-.3)\n",
        "\n",
        "  ax3 = fig.add_subplot(spec[1, 2])\n",
        "  plotting.plot_tsne_by_round(ax3, Data.tsne, round_idx=2)\n",
        "  ax3.set_title('Round 3', y=-.3)\n",
        "\n",
        "\n",
        "make_tsne_figure()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aXokwfP-aFrp"
      },
      "source": [
        "## Distance by log KD"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qb9XumLa6iK1"
      },
      "outputs": [],
      "source": [
        "def _plot_log_kd_by_distance(ax, df, n):\n",
        "  df = df[(df['round'] == 2) | (df['source_num_mutations'] == 0)]\n",
        "  plotting.plot_log_kd_by_distance(\n",
        "      ax=ax, df=df, target_name=COV2_WT, n=n)\n",
        "  ax.set_title(f'Top {n} sequences in round 3', y=-.3)\n",
        "\n",
        "_, ax = plt.subplots(figsize=(7, 5))\n",
        "_plot_log_kd_by_distance(ax, Data.unified, n=100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K0ZivfmTtoyR"
      },
      "outputs": [],
      "source": [
        "def _plot_log_kd_by_mpd(ax, df, n: int):\n",
        "  df = df[df['source_design'] != SHUFFLED]\n",
        "  plotting.plot_log_kd_by_mean_pairwise_distance(\n",
        "      ax, df, n=n, target_name=COV2_WT)\n",
        "\n",
        "_, ax = plt.subplots(figsize=(7, 5))\n",
        "_plot_log_kd_by_mpd(ax, Data.unified.query('round == 2'), n=50)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bPWxu43bO6Ls"
      },
      "source": [
        "## Dendogram"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dpYcZlj6xjWx"
      },
      "outputs": [],
      "source": [
        "_, ax = plt.subplots(figsize=(5, 12))\n",
        "plotting.plot_dendogram(ax, Data.bli_v2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Qf0HH_cbZDP"
      },
      "source": [
        "# BLI and neutralization experiments"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cz7UdzML6o4C"
      },
      "outputs": [],
      "source": [
        "def join_curves_with_bli_data(\n",
        "    curves: pd.DataFrame, bli: pd.DataFrame\n",
        ") -\u003e pd.DataFrame:\n",
        "  \"\"\"Joins BLI curves with BLI summary statistics.\"\"\"\n",
        "  return helper.safe_merge(curves, bli, on=['target_name', 'source_pid'])\n",
        "\n",
        "\n",
        "def aggregate_curves_over_replicas(curves: pd.DataFrame) -\u003e pd.DataFrame:\n",
        "  \"\"\"Aggregates BLI or neutralization curves over replicas.\"\"\"\n",
        "  if 'meas_replica' in curves.columns:\n",
        "    # Neutralization columns\n",
        "    by = ['target_name', 'source_pid', 'meas_replica', 'conc']\n",
        "  else:\n",
        "    # BLI columns\n",
        "    by = ['target_name', 'source_pid', 'conc', 'time']\n",
        "  return (\n",
        "      curves.drop(columns=['replica', 'source_key'], errors='ignore')\n",
        "      .groupby(by)\n",
        "      .mean()\n",
        "      .reset_index()\n",
        "  )\n",
        "\n",
        "\n",
        "def preprocess_curves(curves: pd.DataFrame, bli: pd.DataFrame) -\u003e pd.DataFrame:\n",
        "  \"\"\"Preprocesses BLI or neutralization curves.\"\"\"\n",
        "  return curves.pipe(aggregate_curves_over_replicas).pipe(\n",
        "      join_curves_with_bli_data, bli=bli\n",
        "  )\n",
        "\n",
        "\n",
        "D.bli_curves = preprocess_curves(\n",
        "    covid.load_df('bli_v2_curves.csv.gz', compression='gzip'), D.bli_v2\n",
        ")\n",
        "D.neutralization_curves = preprocess_curves(\n",
        "    covid.load_df('bli_v2_neutralization_curves.csv.gz', compression='gzip'),\n",
        "    D.bli_v2,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i1c1pNqFbR5h"
      },
      "source": [
        "## BLI charts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hzOOTafvRgC5"
      },
      "outputs": [],
      "source": [
        "BLI_METRIC_MAPPING = immutabledict.immutabledict({\n",
        "    'kd': 'KD (nM)',\n",
        "    'ic50': 'IC50 (nM)',\n",
        "    'prod': 'Expression (mg/L))',\n",
        "})\n",
        "\n",
        "\n",
        "BLI_COLOR_PALETTES = immutabledict.immutabledict({\n",
        "    'kd': 'spectral_a',\n",
        "    'ic50': 'plasma_d',\n",
        "    'prod': 'viridis',\n",
        "})\n",
        "\n",
        "\n",
        "BLI_LABELS_MAPPING = collections.OrderedDict([\n",
        "    (63207, 'Seq1'),\n",
        "    (63411, 'Seq2'),\n",
        "    (64284, 'Seq3'),\n",
        "    (63052, 'Seq4'),\n",
        "    (61433, 'Seq5'),\n",
        "    (62541, 'Seq6'),\n",
        "    (61300, 'Seq7'),\n",
        "    (63385, 'Seq8'),\n",
        "    (60461, 'Seq9'),\n",
        "    (60441, 'Seq10'),\n",
        "    (62851, 'Seq11'),\n",
        "    (9643, 'Seq12'),\n",
        "    (60008, 'VHH-72'),\n",
        "])\n",
        "\n",
        "\n",
        "IC50_LOWER_BOUNDS = (500, 1000)\n",
        "\n",
        "PARENT_COLOR = '#636363'\n",
        "\n",
        "\n",
        "def bli_value_is_lower_bound(metric: str, record: pd.Series) -\u003e bool:\n",
        "  return metric == 'ic50' and record[metric] in IC50_LOWER_BOUNDS\n",
        "\n",
        "\n",
        "def _plot_bli_bar(\n",
        "    data,\n",
        "    metric,\n",
        "    max_value,\n",
        "    labels,\n",
        "    show_labels,\n",
        "    width=70,\n",
        "    height=210):\n",
        "  \"\"\"Plots the BLI measurement of a single target (one column).\"\"\"\n",
        "  df = data.assign(\n",
        "      rank_metric=lambda df: df[f'{metric}'].rank(method='first'),\n",
        "      is_best=lambda df: metric != 'prod' and df['rank_metric'] == 1)\n",
        "\n",
        "  parent_record = df.query('source_is_parent')\n",
        "  parent_record = parent_record.iloc[0]\n",
        "  parent_is_lower_bound = bli_value_is_lower_bound(metric, parent_record)\n",
        "  metric_fold = f'{metric}_fold'\n",
        "\n",
        "  def _get_text_label(row):\n",
        "    value = row[metric]\n",
        "    if value \u003e max_value:\n",
        "        label = f'\u003e{int(max_value):d}'\n",
        "    elif bli_value_is_lower_bound(metric, row):\n",
        "      label = f'\u003e{int(value):d}'\n",
        "    else:\n",
        "      label = f'{value:.1f}'\n",
        "    if row['is_best'] and metric_fold in row:\n",
        "      fold_value = row[metric_fold]\n",
        "      fold_label = f'{fold_value:.1f}x'\n",
        "      if parent_is_lower_bound:\n",
        "        fold_label = f'\u003e{fold_label}'\n",
        "      label = f'{label} ({fold_label})'\n",
        "    return label\n",
        "\n",
        "  df = df.assign(text=lambda df: df.apply(_get_text_label, axis=1))\n",
        "\n",
        "  tooltip = helper.ordered_intersection(\n",
        "      ['key', 'source_std_group', 'source_num_mutations',\n",
        "       'kd', 'kd_fold', 'ic50', 'ic50_fold', 'ka', 'kdis'],\n",
        "      df.columns)\n",
        "\n",
        "  def _get_y(axis=None):\n",
        "    return alt.Y(\n",
        "        'key:N',\n",
        "        title='',\n",
        "        sort=list(labels.values()),\n",
        "        axis=axis)\n",
        "\n",
        "  base = alt.Chart(df).encode(\n",
        "      y=_get_y(axis=alt.Axis(labels=False)),\n",
        "      x=alt.X(metric, title=BLI_METRIC_MAPPING.get(metric, metric)),\n",
        "      tooltip=tooltip,\n",
        "  )\n",
        "\n",
        "  bars = base.mark_bar().encode(\n",
        "      x=alt.X(metric, scale=alt.Scale(domain=(0, max_value), clamp=True)),\n",
        "      color=alt.condition(\n",
        "          alt.datum.source_is_parent,\n",
        "          alt.value(PARENT_COLOR),\n",
        "          plotting.get_alt_color(\n",
        "            metric,\n",
        "            palette=BLI_COLOR_PALETTES.get(metric),\n",
        "            legend=alt.Legend(\n",
        "                orient='none',\n",
        "                direction='horizontal',\n",
        "                legendX=0,\n",
        "                legendY=-50),\n",
        "            title=helper.get_unique_value(df['target_name'])))\n",
        "  )\n",
        "\n",
        "  def _get_labels(condition, color='black', **kwargs):\n",
        "    return alt.Chart(df).mark_text(\n",
        "        align='right',\n",
        "        **kwargs\n",
        "    ).encode(\n",
        "        x=alt.value(-5),\n",
        "        y=_get_y(),\n",
        "        text='key:N',\n",
        "        color=alt.condition(condition, alt.value(color), alt.value('black')),\n",
        "        opacity=alt.condition(\n",
        "            condition, alt.value(1), alt.value(0)\n",
        "        ) if show_labels else alt.value(0)\n",
        "    )\n",
        "\n",
        "  def _get_text(condition, **kwargs):\n",
        "    return base.mark_text(\n",
        "        align='left',\n",
        "        baseline='middle',\n",
        "        dx=3,\n",
        "        **kwargs\n",
        "    ).encode(\n",
        "        text='text',\n",
        "        opacity=alt.condition(condition, alt.value(1), alt.value(0))\n",
        "    )\n",
        "\n",
        "  chart = (\n",
        "      bars\n",
        "      + _get_text(~alt.datum.is_best)\n",
        "      + _get_text(alt.datum.is_best, color='red', fontWeight='bold')\n",
        "  )\n",
        "  if show_labels:\n",
        "    chart += (\n",
        "      _get_labels(~alt.datum.source_is_parent, color='black')\n",
        "      + _get_labels(\n",
        "          alt.datum.source_is_parent, color=PARENT_COLOR, fontWeight='bold'))\n",
        "  return chart.properties(width=width, height=height)\n",
        "\n",
        "\n",
        "def plot_bli_bars(bli, metric, max_value, width=70, **kwargs):\n",
        "  \"\"\"Plots the BLI measurement of multiple target.\"\"\"\n",
        "  label_by_key = BLI_LABELS_MAPPING\n",
        "  df = (\n",
        "      bli\n",
        "      .assign(key=lambda df: df['source_key'].map(label_by_key)))\n",
        "\n",
        "  charts = []\n",
        "  target_names = helper.ordered_intersection(\n",
        "      TARGET_NAME_ORDER, df['target_name'])\n",
        "  num_cols = 0\n",
        "  for target_name in target_names:\n",
        "    df_target = df.query('target_name == @target_name')\n",
        "    if df_target[metric].isna().all():\n",
        "      continue\n",
        "    num_cols += 1\n",
        "    charts.append(\n",
        "        _plot_bli_bar(\n",
        "            df_target,\n",
        "            labels=label_by_key,\n",
        "            show_labels=num_cols == 1,\n",
        "            metric=metric,\n",
        "            max_value=max_value,\n",
        "            width=width,\n",
        "            **kwargs))\n",
        "  return (\n",
        "      alt.hconcat(*charts)\n",
        "      .resolve_scale(color='independent')\n",
        "      .configure_legend(gradientLength=width))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TU5HnZB7Ahkt"
      },
      "outputs": [],
      "source": [
        "plot_bli_bars(\n",
        "    helper.map_columns(D.bli_v2, target_name=TARGET_NAME_MAPPING),\n",
        "    metric='kd',\n",
        "    max_value=50,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BG6yO8iebrnV"
      },
      "outputs": [],
      "source": [
        "plot_bli_bars(\n",
        "    helper.map_columns(D.bli_v2, target_name=TARGET_NAME_MAPPING),\n",
        "    metric='ic50',\n",
        "    max_value=500,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U5gXmV9ox7nv"
      },
      "source": [
        "## BLI curves"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hSA32cyh6uAg"
      },
      "outputs": [],
      "source": [
        "def plot_bli_curves(data, labels, with_legend=True, figsize=(6, 4)):\n",
        "  _, ax = plt.subplots(figsize=figsize)\n",
        "  sns.lineplot(\n",
        "      data=data,\n",
        "      x='time',\n",
        "      y='data',\n",
        "      hue='label',\n",
        "      hue_order=labels,\n",
        "      palette='tab10',\n",
        "      linewidth=3,\n",
        "      alpha=0.8,\n",
        "      ax=ax,\n",
        "  )\n",
        "  ax.set(xlabel='Time (s)', ylabel='Response (nm)')\n",
        "  if with_legend:\n",
        "    ax.legend(title='', ncol=1, bbox_to_anchor=(1.0, 1.0))\n",
        "  else:\n",
        "    ax.legend_.remove()\n",
        "  return ax\n",
        "\n",
        "\n",
        "def plot_bli_curves_facet(data, labels, concentation=30):\n",
        "  df = data.query('conc == @concentation and label in @labels').pipe(\n",
        "      helper.map_columns, target_name=TARGET_NAME_MAPPING\n",
        "  )\n",
        "  target_names = helper.ordered_intersection(\n",
        "      TARGET_NAME_ORDER, df['target_name']\n",
        "  )\n",
        "  for target_name in target_names:\n",
        "    plot_bli_curves(\n",
        "        data=df.query('target_name == @target_name'),\n",
        "        labels=labels,\n",
        "        with_legend=target_name == target_names[-1],\n",
        "    ).set_title(target_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QjrRf7LqAyQD"
      },
      "outputs": [],
      "source": [
        "plot_bli_curves_facet(D.bli_curves, labels=['Seq1', 'Seq2', 'Seq3', 'VHH-72'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "THGYj8AAPCR9"
      },
      "source": [
        "## Neutralization curves"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ku-w0d3mgPIf"
      },
      "outputs": [],
      "source": [
        "def plot_neutralization_curves(data, labels, with_legend=True, figsize=(6, 4)):\n",
        "  _, ax = plt.subplots(figsize=figsize)\n",
        "  sns.lineplot(\n",
        "      data=data,\n",
        "      x='conc',\n",
        "      y='value',\n",
        "      hue='label',\n",
        "      hue_order=labels,\n",
        "      palette='tab10',\n",
        "      linewidth=3,\n",
        "      marker='o',\n",
        "      dashes=True,\n",
        "      err_style='bars',\n",
        "      errorbar='sd',\n",
        "      ax=ax,\n",
        "  )\n",
        "  ax.set(xlabel='Concentation (nM)', xscale='log', ylabel='% Infectivity')\n",
        "  if with_legend:\n",
        "    ax.legend(title='', ncol=1, bbox_to_anchor=(1.0, 1.0))\n",
        "  else:\n",
        "    ax.legend_.remove()\n",
        "  return ax\n",
        "\n",
        "\n",
        "def plot_neutralization_curves_facet(data, labels):\n",
        "  df = data.query('label in @labels').pipe(\n",
        "      helper.map_columns, target_name=TARGET_NAME_MAPPING\n",
        "  )\n",
        "  target_names = helper.ordered_intersection(\n",
        "      TARGET_NAME_ORDER, df['target_name']\n",
        "  )\n",
        "  for target_name in target_names:\n",
        "    plot_neutralization_curves(\n",
        "        data=df.query('target_name == @target_name'),\n",
        "        labels=labels,\n",
        "        with_legend=target_name == target_names[-1],\n",
        "    ).set_title(target_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8gib9g-SGn-n"
      },
      "outputs": [],
      "source": [
        "plot_neutralization_curves_facet(\n",
        "    D.neutralization_curves, labels=['Seq1', 'Seq2', 'Seq3', 'VHH-72']\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QTP60iFFQA3v"
      },
      "source": [
        "# Model evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VtRpSGZRROLC"
      },
      "source": [
        "## Constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T2QO_9GdRZkq"
      },
      "outputs": [],
      "source": [
        "# Datasets that were used for training.\n",
        "CONFIG_DSET_MAPPING = immutabledict.immutabledict({\n",
        "    'r1': 'Round 1',\n",
        "    'r2': 'Round 1+2',\n",
        "    'r3': 'Round 1+2+3',\n",
        "})\n",
        "\n",
        "CONFIG_DSET_ORDER = tuple(CONFIG_DSET_MAPPING.values())\n",
        "\n",
        "CONFIG_MODEL_MAPPING = immutabledict.immutabledict({\n",
        "    'cnn': 'CNN',\n",
        "    'lgb': 'LGB',\n",
        "    'linear': 'Baseline',\n",
        "})\n",
        "\n",
        "MODEL_ORDER =  ('Baseline', 'CNN', 'LGB')\n",
        "\n",
        "METRIC_MAPPING = immutabledict.immutabledict({\n",
        "    'spearmanr': 'Spearman R',\n",
        "    'roc_auc_score_parent': 'ROC AUC VHH-72',\n",
        "})\n",
        "\n",
        "MODEL_COLUMN_MAPPING = immutabledict.immutabledict({\n",
        "    'config_dset': 'Training set',\n",
        "    'config_model': 'Model',\n",
        "})\n",
        "\n",
        "\n",
        "# Targets that were used to train both r1 and r2 models.\n",
        "JOINT_MODEL_TARGETS = (\n",
        "    'SARS-CoV1_RBD',\n",
        "    'SARS-CoV2_RBD',\n",
        "    'SARS-CoV2_RBD_G502D',\n",
        "    'SARS-CoV2_RBD_N439K',\n",
        "    'SARS-CoV2_RBD_N501D',\n",
        "    'SARS-CoV2_RBD_N501F',\n",
        "    'SARS-CoV2_RBD_S477N',\n",
        "    'SARS-CoV2_RBD_V367F',\n",
        ")\n",
        "\n",
        "def map_model_columns(df, ignore_missing=True, **kwargs):\n",
        "  return helper.map_columns(\n",
        "      df,\n",
        "      config_model=CONFIG_MODEL_MAPPING,\n",
        "      metric=METRIC_MAPPING,\n",
        "      config_dset=CONFIG_DSET_MAPPING,\n",
        "      **kwargs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yfGiYQEAVlDp"
      },
      "source": [
        "## Read data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zcwrMwJOSYaL"
      },
      "outputs": [],
      "source": [
        "def _read_model_eval_df(name):\n",
        "  filename = os.path.join(covid.DATA_DIR, 'model_scores', f'{name}.csv.gz')\n",
        "  helper.tprint(filename)\n",
        "  return helper.read_csv(filename)\n",
        "\n",
        "\n",
        "M = helper.Bunch()\n",
        "# Hold-out model evaluation scores.\n",
        "M.scores = _read_model_eval_df('scores')\n",
        "# Hold-out model evaluation scores per distance.\n",
        "M.scores_dist = _read_model_eval_df('scores_dist')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BzHVFXyvSqcU"
      },
      "outputs": [],
      "source": [
        "def _read_cv_scores():\n",
        "  name_by_config_dset = dict(\n",
        "      r1='cv_scores_r0',\n",
        "      r2='cv_scores_r1',\n",
        "      r3='cv_scores_r2',\n",
        "  )\n",
        "  return (\n",
        "      pd.concat([\n",
        "          _read_model_eval_df(name).assign(config_dset=config_dset)\n",
        "          for config_dset, name in name_by_config_dset.items()\n",
        "          ], ignore_index=True)\n",
        "      .rename(columns={'config': 'config_model'})\n",
        "  )\n",
        "\n",
        "# Cross-validation scores.\n",
        "M.cv_scores = _read_cv_scores()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KuPfezyLoKqQ"
      },
      "source": [
        "## Per round"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qZF8qiDCS5Tj"
      },
      "outputs": [],
      "source": [
        "def plot_model_performance_per_round(\n",
        "    scores,\n",
        "    metric,\n",
        "    metric_label,\n",
        "    config_dset=None,\n",
        "    ylim=None,\n",
        "    figsize=(5, 5),\n",
        "    **kwargs):\n",
        "  df = scores.query(\n",
        "      'metric == @metric and target_name in @JOINT_MODEL_TARGETS and '\n",
        "      'config_model in @CONFIG_MODEL_MAPPING')\n",
        "  if 'source_dset' in df.columns:\n",
        "    # The dataset that was used for evaluation. \"r2\" corresponds to the 3rd\n",
        "    # library, which was not used for model-training.\n",
        "    df = df.query('source_dset == \"r2\"')\n",
        "  if config_dset:\n",
        "    df = df.query('config_dset == @config_dset')\n",
        "  df = map_model_columns(df)\n",
        "\n",
        "  _, ax = plt.subplots(figsize=figsize)\n",
        "  ax = sns.barplot(\n",
        "      data=df,\n",
        "      x='config_dset',\n",
        "      order=helper.ordered_intersection(CONFIG_DSET_ORDER, df['config_dset']),\n",
        "      y='value',\n",
        "      hue='config_model',\n",
        "      palette=COLOR_PALETTE,\n",
        "      hue_order=MODEL_ORDER,\n",
        "      ax=ax,\n",
        "      **kwargs)\n",
        "  ax.set_xlabel('Training set')\n",
        "  ax.set_ylabel(metric_label)\n",
        "  if ylim is not None:\n",
        "    ax.set_ylim(ylim)\n",
        "  ax.legend(bbox_to_anchor=(.5, 1.1), loc='center', ncol=3, frameon=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S1ID4kfZdDIw"
      },
      "source": [
        "### Hold-out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sarUlgxJHZH9"
      },
      "outputs": [],
      "source": [
        "plot_model_performance_per_round(\n",
        "    # Plot the performance of the joint regressor-classifier model (\"joint\").\n",
        "    # config_type == \"reg\"/\"cla\" selects the results of the\n",
        "    # regressor/classififer.\n",
        "    M.scores.query('config_type == \"joint\"'),\n",
        "    metric='spearmanr',\n",
        "    metric_label='Hold-out Spearman R')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kduVXz-xbL-R"
      },
      "outputs": [],
      "source": [
        "plot_model_performance_per_round(\n",
        "    M.scores.query('config_type == \"joint\"'),\n",
        "    metric='roc_auc_score_parent',\n",
        "    metric_label='Hold-out ROC AUC VHH-72',\n",
        "    ylim=(0.5, None))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3BI-CmDka2MK"
      },
      "source": [
        "### Cross-validation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N-i6gCalHlca"
      },
      "outputs": [],
      "source": [
        "plot_model_performance_per_round(\n",
        "    M.cv_scores,\n",
        "    metric='spearmanr',\n",
        "    metric_label='CV Spearman R',\n",
        "    config_dset=['r1', 'r2'],\n",
        "    ylim=(0.5, None))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g1LzN85HHnFd"
      },
      "outputs": [],
      "source": [
        "plot_model_performance_per_round(\n",
        "    M.cv_scores,\n",
        "    metric='roc_auc_score_parent',\n",
        "    metric_label='CV ROC AUC VHH-72',\n",
        "    config_dset=['r1', 'r2'],\n",
        "    ylim=(0.75, None))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R0Nzzfhlb-Ug"
      },
      "source": [
        "## Per distance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zCTXgyoooIAT"
      },
      "outputs": [],
      "source": [
        "def plot_model_performance_per_distance(\n",
        "    all_scores,\n",
        "    metric='spearmanr',\n",
        "    with_num_samples=True,\n",
        "    ylim=None,\n",
        "    ylabel=None,\n",
        "    ylabel_prefix=None,\n",
        "    figsize=(12, 5.5),\n",
        "    **kwargs):\n",
        "  scores = all_scores.query(\n",
        "      'target_name in @JOINT_MODEL_TARGETS and config_type == \"joint\" and '\n",
        "      'config_model in @CONFIG_MODEL_MAPPING and source_dset == \"r2\"')\n",
        "  df = (\n",
        "      scores\n",
        "      .query('metric == @metric')\n",
        "      .pipe(map_model_columns)\n",
        "      .rename(columns=MODEL_COLUMN_MAPPING))\n",
        "\n",
        "  if with_num_samples:\n",
        "    # Plot the number of samples that were used for the evaluation.\n",
        "    evaluator = helper.get_unique_value(df['evaluator'])\n",
        "    df_samples = (\n",
        "        scores.query('metric == \"samples\" and evaluator == @evaluator'))\n",
        "    display.display(\n",
        "        df_samples.pivot_table(\n",
        "            index=['config_dset', 'config_model'],\n",
        "            columns='source_num_mutations',\n",
        "            values='value'))\n",
        "    samples_per_dist = (\n",
        "        df_samples.groupby('source_num_mutations')['value'].mean()\n",
        "        .astype('int').to_dict())\n",
        "\n",
        "    def _get_xlabel(dist):\n",
        "      n = samples_per_dist[dist]\n",
        "      if n \u003e 1000:\n",
        "        n = f'{int(n // 1000)}k'\n",
        "      return f'{dist:d}\\nn={n}'\n",
        "\n",
        "    df['x'] = df['source_num_mutations'].apply(_get_xlabel)\n",
        "  else:\n",
        "    df['x'] = df['source_num_mutations']\n",
        "  order = sorted(set(df['x']))\n",
        "\n",
        "  _, ax = plt.subplots(figsize=figsize)\n",
        "  sns.lineplot(\n",
        "      data=df,\n",
        "      x='x',\n",
        "      y='value',\n",
        "      hue='Model',\n",
        "      hue_order=MODEL_ORDER,\n",
        "      palette=dict(COLOR_PALETTE),\n",
        "      style='Training set',\n",
        "      style_order=['Round 1+2', 'Round 1'],\n",
        "      markers=True,\n",
        "      dashes=True,\n",
        "      markersize=10,\n",
        "      err_style='bars',\n",
        "      ax=ax)\n",
        "  ax.tick_params(axis='x', which='major', labelsize=15)\n",
        "  ax.set_xlabel(f'Number of mutations from {PARENT_NAME}')\n",
        "\n",
        "  if ylabel is None:\n",
        "    ylabel = helper.get_unique_value(df['metric'])\n",
        "  if ylabel_prefix is not None:\n",
        "    ylabel = f'{ylabel_prefix} {ylabel}'\n",
        "  ax.set_ylabel(ylabel)\n",
        "  ax.legend(\n",
        "      loc='lower left', ncol=2, frameon=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UBcf6ND9HpjJ"
      },
      "outputs": [],
      "source": [
        "plot_model_performance_per_distance(\n",
        "    M.scores_dist,\n",
        "    metric='spearmanr',\n",
        "    ylabel_prefix='Hold-out',\n",
        "    ylim=(0.2, None))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HWt4svSten82"
      },
      "outputs": [],
      "source": [
        "plot_model_performance_per_distance(\n",
        "    M.scores_dist,\n",
        "    metric='roc_auc_score_parent',\n",
        "    ylabel_prefix='Hold-out',\n",
        "    ylim=(0.2, None))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jjpfhdx-KCPG"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
